import json
import torch
from transformers import AutoTokenizer
from data_loader import get_dataloader
from model import ChineseMedicalBertQA
from trainer import QATrainer

from transformers import AutoTokenizer, AutoModelForQuestionAnswering

def train():
    # Set device
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f'Using device: {device}')
    
    # 定义本地模型路径
    model_path = 'D:/TCM-ai/models/bert-base-chinese'
    
    # Initialize tokenizer and model from local files
    print('Initializing tokenizer and model from local files...')
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForQuestionAnswering.from_pretrained(model_path)
    
    # Load datasets
    print('Loading datasets...')
    try:
        train_loader = get_dataloader(r'D:/TCM-ai/data/round1_train_0907.json', tokenizer, batch_size=32)
        print(f'Loaded {len(train_loader.dataset)} training examples.')
    except Exception as e:
        print(f'Error loading dataset: {e}')
        return
    
    # Train model
    print('Starting training...')
    trainer = QATrainer(model, train_loader, device=device)
    trainer.train()
    
    print('Training completed!')
def predict():
    # Set device
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # Initialize tokenizer
    tokenizer = AutoTokenizer.from_pretrained('bert-base-chinese')
    
    # Initialize model
    model = ChineseMedicalBertQA()
    
    # Create predictor (this would need to be implemented)
    # predictor = Predictor(model, tokenizer, '../user_data/model_data/final_model.pth', device=device)
    
    # Predict on test set
    print('Generating predictions for test set...')
    # predictor.predict_file(
    #     '../data/round1_test_0907.json', 
    #     '../prediction_result/result.json'
    # )
    
    print('Prediction completed!')

if __name__ == '__main__':
    import sys
    
    if len(sys.argv) != 2:
        print("Usage: python main.py <train|predict>")
        sys.exit(1)
    
    mode = sys.argv[1]
    
    if mode == 'train':
        train()
    elif mode == 'predict':
        predict()
    else:
        print("Invalid mode. Use 'train' or 'predict'")
        sys.exit(1)